package cn.jiangzeyin.aes;

import cn.jiangzeyin.aes.util.Base64Decoder;
import cn.jiangzeyin.aes.util.Base64Encoder;

import javax.crypto.*;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.PBEKeySpec;
import javax.crypto.spec.SecretKeySpec;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.spec.InvalidKeySpecException;

public class AES {

//    public static String PASSWORD = "";

    private final int HASH_ITERATIONS = 10000;
    private final int KEY_LENGTH = 256;
    private final String CIPHERMODEPADDING = "AES/CBC/PKCS5Padding";
//    private final String ECB = "AES/ECB/PKCS7Padding";

    private char[] humanPassphrase = {'P', 'e', 'r', ' ', 'v', 'a', 'l', 'l', 'u', 'm', ' ', 'd', 'u', 'c', 'e', 's', ' ', 'L', 'a', 'b', 'a', 'n', 't'};

    private byte[] salt = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0xA, 0xB, 0xC, 0xD, 0xE, 0xF}; // must save this for next time we want the key

    private PBEKeySpec myKeyspec = new PBEKeySpec(humanPassphrase, salt, HASH_ITERATIONS, KEY_LENGTH);


    private SecretKey sk = null;
    private SecretKeySpec skforAES = null;
    private byte[] iv = {0xA, 1, 0xB, 5, 4, 0xF, 7, 9, 0x17, 3, 1, 6, 8, 0xC, 0xD, 91};

    private IvParameterSpec IV;

    private void initPassword(String password) {
        try {
            if (password != null && !password.isEmpty()) {
                myKeyspec = new PBEKeySpec(password.toCharArray(), salt,
                        HASH_ITERATIONS, KEY_LENGTH);
            }
            String KEY_GENERATION_ALG = "PBKDF2WithHmacSHA1";
            SecretKeyFactory keyfactory = SecretKeyFactory.getInstance(KEY_GENERATION_ALG);
            sk = keyfactory.generateSecret(myKeyspec);
        } catch (NoSuchAlgorithmException | InvalidKeySpecException nsae) {
            nsae.printStackTrace();
        }

        // This is our secret key. We could just save this to a file instead of
        // regenerating it
        // each time it is needed. But that file cannot be on the device (too
        // insecure). It could
        // be secure if we kept it on a server accessible through https.
        byte[] skAsByteArray = sk.getEncoded();
        skforAES = new SecretKeySpec(skAsByteArray, "AES");
        IV = new IvParameterSpec(iv);
    }

//    public String aesEncryptECB(String password, byte[] plaintext) {
//        byte[] result = null;
//        try {
//            SecretKeyFactory factory = SecretKeyFactory
//                    .getInstance("PBKDF2WithHmacSHA1");
//            KeySpec spec = new PBEKeySpec(password.toCharArray(), salt, 1024,
//                    128);
//            SecretKey tmp = factory.generateSecret(spec);
//            SecretKey secret = new SecretKeySpec(tmp.getEncoded(), ECB);
//            Cipher cipher = Cipher.getInstance(ECB);
//            cipher.init(Cipher.ENCRYPT_MODE, secret);
//            result = cipher.doFinal(plaintext);
//        } catch (InvalidKeyException | InvalidKeySpecException | BadPaddingException | IllegalBlockSizeException | NoSuchPaddingException | NoSuchAlgorithmException e) {
//            e.printStackTrace();
//        }
//        return Base64Encoder.encode(result);
//    }

//    private static String getDateToString(long time) {
//        Date d = new Date(time);
//        SimpleDateFormat sf = new SimpleDateFormat("yyyyMMddHHmmss");
//        return sf.format(d);
//    }

    public String encrypt(String password, byte[] plaintext) {
        initPassword(password);
        byte[] ciphertext = encrypt(CIPHERMODEPADDING, skforAES, IV, plaintext);
//        return Base64.getEncoder().encodeToString(ciphertext);
        if (ciphertext != null) {
            return Base64Encoder.encode(ciphertext);
        }
        return null;
    }

    public String decrypt(String password, String ciphertext_base64) {
        initPassword(password);

//        byte[] s = Base64.getDecoder().decode(ciphertext_base64);
        byte[] s = Base64Decoder.decodeToBytes(ciphertext_base64);
        byte[] d = decrypt(CIPHERMODEPADDING, skforAES, IV, s);
        if (d == null)
            return null;
        return new String(d);
    }

//    public String aesDecryptECB(String password, String ciphertext_base64) {
//        try {
//            SecretKeyFactory factory = SecretKeyFactory
//                    .getInstance("PBKDF2WithHmacSHA1");
//            KeySpec spec = new PBEKeySpec(password.toCharArray(), salt, 1024,
//                    128);
//            SecretKey tmp = factory.generateSecret(spec);
//            SecretKey secret = new SecretKeySpec(tmp.getEncoded(), ECB);
//            Cipher cipher = Cipher.getInstance(ECB);
//            cipher.init(Cipher.DECRYPT_MODE, secret);
//            byte[] result = cipher.doFinal(Base64Decoder.decodeToBytes(ciphertext_base64));
//            return new String(result, "UTF-8");
//        } catch (NoSuchAlgorithmException | NoSuchPaddingException | BadPaddingException | UnsupportedEncodingException | InvalidKeySpecException | InvalidKeyException | IllegalBlockSizeException e) {
//            e.printStackTrace();
//        }
//        return null;
//    }

    // Use this method if you want to add the padding manually
    // AES deals with messages in blocks of 16 bytes.
    // This method looks at the length of the message, and adds bytes at the end
    // so that the entire message is a multiple of 16 bytes.
    // the padding is a series of bytes, each set to the total bytes added (a
    // number in range 1..16).
//    public byte[] addPadding(byte[] plain) {
//        byte plainpad[] = null;
//        int shortage = 16 - (plain.length % 16);
//        // if already an exact multiple of 16, need to add another block of 16
//        // bytes
//        if (shortage == 0)
//            shortage = 16;
//
//        // reallocate array bigger to be exact multiple, adding shortage bits.
//        plainpad = new byte[plain.length + shortage];
//        for (int i = 0; i < plain.length; i++) {
//            plainpad[i] = plain[i];
//        }
//        for (int i = plain.length; i < plain.length + shortage; i++) {
//            plainpad[i] = (byte) shortage;
//        }
//        return plainpad;
//    }
//
//    // Use this method if you want to remove the padding manually
//    // This method removes the padding bytes
//    public byte[] dropPadding(byte[] plainpad) {
//        byte plain[] = null;
//        int drop = plainpad[plainpad.length - 1]; // last byte gives number of
//        // bytes to drop
//
//        // reallocate array smaller, dropping the pad bytes.
//        plain = new byte[plainpad.length - drop];
//        for (int i = 0; i < plain.length; i++) {
//            plain[i] = plainpad[i];
//            plainpad[i] = 0; // don't keep a copy of the decrypt
//        }
//        return plain;
//    }

    private byte[] encrypt(String cmp, SecretKey sk, IvParameterSpec IV, byte[] msg) {
        try {
            Cipher c = Cipher.getInstance(cmp);
            c.init(Cipher.ENCRYPT_MODE, sk, IV);
            return c.doFinal(msg);
        } catch (NoSuchAlgorithmException | NoSuchPaddingException | InvalidAlgorithmParameterException | InvalidKeyException | BadPaddingException | IllegalBlockSizeException nsae) {
            nsae.printStackTrace();
        }
        return null;
    }

//    public byte[] encryptECB(String cmp, SecretKey sk, byte[] msg) {
//        try {
//            Cipher c = Cipher.getInstance(cmp);
//            c.init(Cipher.ENCRYPT_MODE, sk);
//            return c.doFinal(msg);
//        } catch (NoSuchAlgorithmException | NoSuchPaddingException | InvalidKeyException | BadPaddingException | IllegalBlockSizeException nsae) {
//            nsae.printStackTrace();
//        }
//        return null;
//    }

    private byte[] decrypt(String cmp, SecretKey sk, IvParameterSpec IV, byte[] ciphertext) {
        try {
            Cipher c = Cipher.getInstance(cmp);
            c.init(Cipher.DECRYPT_MODE, sk, IV);
            return c.doFinal(ciphertext);
        } catch (NoSuchAlgorithmException | BadPaddingException | IllegalBlockSizeException | InvalidAlgorithmParameterException | InvalidKeyException | NoSuchPaddingException nsae) {
            nsae.printStackTrace();
        }
        return null;
    }

//    public byte[] decryptECB(String cmp, SecretKey sk, byte[] ciphertext) {
//        try {
//            Cipher c = Cipher.getInstance(cmp);
//            c.init(Cipher.DECRYPT_MODE, sk);
//            return c.doFinal(ciphertext);
//        } catch (NoSuchAlgorithmException | NoSuchPaddingException | InvalidKeyException | BadPaddingException | IllegalBlockSizeException nsae) {
//            nsae.printStackTrace();
//        }
//        return null;
//    }

}